import numpy as np
import math
import cv2
import json
from PIL import Image, ImageDraw, ImageFont


class PoseLabelingUtils:
    def draw_action_label(self, annotated_image, actions, label_position):
        pil_image = Image.fromarray(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image)
        font = ImageFont.truetype("simhei.ttf", 30)
        # label_position = (10, 30)
        text_position = (label_position[0], label_position[1])
        # print(text_position)
        draw.text(text_position, f"{actions}", font=font, fill=(0, 0, 255))
        image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        return image

    def load_keypoints_and_thresholds(self, file_path):
        with open(file_path, 'r') as file:
            data = json.load(file)
            keypoint_index = data['KEYPOINT_INDEXES']
            pose_landmark = data['POSE_LANDMARK']
            face_thresholds = data['FACE_THRESHOLDS']
            arm_thresholds = data['ARM_THRESHOLDS']
            leg_thresholds = data['LEG_THRESHOLDS']
        return keypoint_index, pose_landmark, face_thresholds, arm_thresholds, leg_thresholds

    def calculate_angle(self, point1, point2, point3):
        """计算三个点形成的夹角"""
        # 计算向量
        vector1 = np.array([point2[0] - point1[0], point2[1] - point1[1]])
        vector2 = np.array([point3[0] - point2[0], point3[1] - point2[1]])

        # 计算向量的模
        norm1 = np.linalg.norm(vector1)
        norm2 = np.linalg.norm(vector2)

        # 计算点积
        dot_product = np.dot(vector1, vector2)

        # 计算夹角的余弦值
        cos_theta = dot_product / (norm1 * norm2)

        # 计算夹角（弧度）
        angle_rad = np.arccos(np.clip(cos_theta, -1.0, 1.0))

        # 转换为角度
        angle_deg = np.degrees(angle_rad)
        return angle_deg

    def label_head_pose(self, keypoints):
        """
        根据人脸关键点的位置关系确定头部的姿势标签。

        参数:
            keypoints: 关键点字典，包含人脸的关键点。

        返回:
            头部的姿势标签，如 "正头" 或 "侧头"。
        """
        left_eye = keypoints['left_eye']
        right_eye = keypoints['right_eye']
        left_ear = keypoints['left_ear']
        right_ear = keypoints['right_ear']
        nose = keypoints['nose']

        # 计算关键点之间的距离
        eye_distance = ((left_eye[0] - right_eye[0]) ** 2 + (left_eye[1] - right_eye[1]) ** 2) ** 0.5
        ear_distance = ((left_ear[0] - right_ear[0]) ** 2 + (left_ear[1] - right_ear[1]) ** 2) ** 0.5

        # 设置阈值（这些阈值需要根据实际情况进行调整）
        EYE_DISTANCE_THRESHOLD = 50  # 眼睛距离阈值
        EAR_DISTANCE_THRESHOLD = 70  # 耳朵距离阈值
        ANGLE_THRESHOLD = 10  # 角度阈值（用于判断歪头）
        VERTICAL_DISPLACEMENT_THRESHOLD = 20  # 垂直位移阈值（用于判断抬头和低头）

        # 计算眼睛和鼻子的垂直位置差异（用于判断抬头和低头）
        eye_midpoint_y = (left_eye[1] + right_eye[1]) / 2
        nose_y = nose[1]
        vertical_displacement = abs(eye_midpoint_y - nose_y)

        # 计算头部倾斜角度（简化方法，仅使用眼睛和耳朵的相对位置）
        eye_slope = (right_eye[1] - left_eye[1]) / (right_eye[0] - left_eye[0]) if right_eye[0] != left_eye[
            0] else float('inf')
        ear_slope = (right_ear[1] - left_ear[1]) / (right_ear[0] - left_ear[0]) if right_ear[0] != left_ear[
            0] else float('inf')
        angle_between_lines = abs(math.atan2(eye_slope, -1) - math.atan2(ear_slope, -1))
        angle_in_degrees = math.degrees(angle_between_lines)

        # 判断头部姿势
        if eye_distance < EYE_DISTANCE_THRESHOLD and ear_distance < EAR_DISTANCE_THRESHOLD:
            if vertical_displacement > VERTICAL_DISPLACEMENT_THRESHOLD:
                if eye_midpoint_y > nose_y:
                    return "侧头且抬头"
                else:
                    return "侧头且低头"
            else:
                return "侧头"
        elif angle_in_degrees > ANGLE_THRESHOLD:
            return "歪头"
        elif vertical_displacement > VERTICAL_DISPLACEMENT_THRESHOLD:
            # 垂直位移较大，可能是抬头或低头
            if eye_midpoint_y > nose_y:
                return "抬头"
            else:
                return "低头"
        else:
            return "正头"

    def calculate_arm_angles(self, keypoints_dict, side):
        """
        计算指定手臂的角度。

        参数:
            keypoints_dict: 关键点字典
            side: 'left' 或 'right'

        返回:
            手臂的角度（以度为单位），如果关键点无效则返回 None
        """
        shoulder_key = f'{side}_shoulder'
        elbow_key = f'{side}_elbow'
        wrist_key = f'{side}_wrist'

        shoulder = keypoints_dict.get(shoulder_key)
        elbow = keypoints_dict.get(elbow_key)
        wrist = keypoints_dict.get(wrist_key)

        return self.calculate_angle(shoulder, elbow, wrist)

    def label_arm_pose(self, keypoints):
        left_shoulder = keypoints['left_shoulder']
        right_shoulder = keypoints['right_shoulder']
        left_elbow = keypoints['left_elbow']
        right_elbow = keypoints['right_elbow']
        left_wrist = keypoints['left_wrist']
        right_wrist = keypoints['right_wrist']

        # 计算左右手臂的角度
        left_arm_angle = self.calculate_arm_angles(keypoints, 'left')
        right_arm_angle = self.calculate_arm_angles(keypoints, 'right')

        # 判断左手臂姿势
        if left_wrist[1] < left_shoulder[1]:  # 腕部低于肩部，可能是直举或弯举
            if left_arm_angle is not None and left_arm_angle <= 30:
                left_label = "左臂直举"
            elif left_arm_angle is not None and 30 < left_arm_angle <= 90:
                left_label = "左臂弯举"
            else:
                left_label = "左臂未知"
        elif left_arm_angle is not None and 150 <= left_arm_angle <= 180:
            left_label = "左手平伸"
        elif left_arm_angle is not None and left_arm_angle > 90:
            left_label = "左手弯曲"
        else:
            left_label = "左手放下"

        # 判断右手臂姿势
        if right_wrist[1] < right_shoulder[1]:  # 腕部低于肩部，可能是直举或弯举
            if right_arm_angle is not None and right_arm_angle <= 30:
                right_label = "右臂直举"
            elif right_arm_angle is not None and 30 < right_arm_angle <= 90:
                right_label = "右臂弯举"
            else:
                right_label = "右臂未知"
        elif right_arm_angle is not None and 150 <= right_arm_angle <= 180:
            right_label = "右手平伸"
        elif right_arm_angle is not None and right_arm_angle > 90:
            right_label = "右手弯曲"
        else:
            right_label = "右手放下"

        # 合并左右的标签
        if left_label[-3:] == right_label[-3:]:
            return f"双{left_label[-3:]}"
        else:
            return f"{left_label},{right_label}"

    def calculate_leg_angles(self, keypoints_dict, side):
        """
        计算指定腿部的角度。

        参数:
            keypoints_dict: 关键点字典
            side: 'left' 或 'right'

        返回:
            腿部的角度（以度为单位），如果关键点无效则返回 None
        """
        hip_key = f'{side}_hip'
        knee_key = f'{side}_knee'
        ankle_key = f'{side}_ankle'

        hip = keypoints_dict.get(hip_key)
        knee = keypoints_dict.get(knee_key)
        ankle = keypoints_dict.get(ankle_key)

        return self.calculate_angle(hip, knee, ankle)

    def label_leg_pose(self, keypoints):
        left_hip = keypoints.get('left_hip')
        right_hip = keypoints.get('right_hip')
        left_knee = keypoints.get('left_knee')
        right_knee = keypoints.get('right_knee')
        left_ankle = keypoints.get('left_ankle')
        right_ankle = keypoints.get('right_ankle')

        if (left_ankle[1] > left_hip[1] and right_ankle[1] > right_hip[1] and
                left_knee[1] > left_hip[1] and right_knee[1] > right_hip[1]):
            return "站立"
        else:
            # 左腿弯曲，右腿站立
            if left_ankle[1] < left_hip[1] or left_knee[1] < left_hip[1]:
                if right_ankle[1] > right_hip[1] and right_knee[1] > right_hip[1]:
                    return "左腿弯曲，右腿站立"
                else:
                    return "双腿弯曲"
            # 右腿弯曲，左腿站立
            elif right_ankle[1] < right_hip[1] or right_knee[1] < right_hip[1]:
                if left_ankle[1] > left_hip[1] and left_knee[1] > left_hip[1]:
                    return "左腿站立，右腿弯曲"
                else:
                    return "双腿弯曲"
            # 其他情况视为蹲下
            else:
                return "蹲下"
